package cn.jbolt.ai.klb.entity;

import cn.hutool.core.util.StrUtil;
import cn.jbolt.ai.klb.FileState;
import cn.jbolt.ai.klb.FileType;
import cn.jbolt.util.JSONUtil;
import com.baomidou.mybatisplus.annotation.*;
import com.fasterxml.jackson.annotation.JsonFormat;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.node.ObjectNode;
import io.swagger.v3.oas.annotations.media.Schema;

import java.io.Serializable;
import java.util.Date;


@TableName("ja_ai_doc")
public class AiDoc implements Serializable {
    private static final long serialVersionUID = 1L;

    @TableId(type = IdType.ASSIGN_ID)
    private Long id;

    @Schema(name = "content", description = "内容")
    @TableField("content")
    private String content;


    @Schema(name = "knowledgeBaseId", description = "知识库ID")
    @TableField("knowledge_base_id")
    private Long knowledgeBaseId;


    @Schema(name = "fileId", description = "文件ID")
    @TableField("file_id")
    private Long fileId;

    @Schema(name = "docId", description = "基于doc进行二次训练")
    @TableField("doc_id")
    private Long docId;

    @Schema(name = "linkId", description = "关联的answerId")
    @TableField("link_id")
    private Long linkId;

    @Schema(name = "linkType", description = "关联的类型")
    @TableField("link_type")
    private String linkType;


    @Schema(name = "dataType", description = "数据类型:customText:自定义文本,filePiece:文件块")
    @TableField("data_type")
    private String dataType;

    @Schema(name = "docType", description = "文件类型,例如:normal:普通文档,QA:QA文档")
    @TableField("doc_type")
    private String docType;


    @Schema(name = "docStatus", description = "文档状态:,100:普通文档未拆分,QA文档未训练,200:普通文档拆分中,QA文档训练中,300:普通文档拆分成功,QA文档训练成功" +
            "400:普通文档拆分失败,QA文档训练失败,500:向量化中,600:向量化成功,700:向量化失败")
    @TableField("doc_status")
    private Integer docStatus;


    @Schema(name = "trainTime", description = "训练时间")
    @JsonFormat(locale = "zh", pattern = "yyyy-MM-dd HH:mm:ss", timezone = "GMT+8")
    @TableField("train_time")
    private Date trainTime;

    @Schema(name = "embedTime", description = "向量化时间")
    @JsonFormat(locale = "zh", pattern = "yyyy-MM-dd HH:mm:ss", timezone = "GMT+8")
    @TableField("embed_time")
    private Date embedTime;

    @Schema(name = "createTime", description = "创建时间")
    @JsonFormat(locale = "zh", pattern = "yyyy-MM-dd HH:mm:ss", timezone = "GMT+8")
    @TableField(value = "create_time", fill = FieldFill.INSERT)
    private Date createTime;


    //非数据库字段
    @TableField(exist = false)
    private String docStatuesName;

    @TableField(exist = false)
    private String aiKnowledgeBaseName;

    @TableField(exist = false)
    private String fileName;

    @TableField(exist = false)
    private Double score;

    @TableField(exist = false)
    private String answer;


    @TableField(exist = false)
    private Integer questionNum;

    public AiDoc() {
    }

    public AiDoc(String content, String answer, Long knowledgeBaseId, Long fileId, String dataType, Integer docStatus, String docType) {
        this.content = content;
//        this.answer = answer;
        this.knowledgeBaseId = knowledgeBaseId;
        this.fileId = fileId;
        this.dataType = dataType;
        this.docStatus = docStatus;
        this.docType = docType;
    }

    /**
     * 转成前端组件需要的JsonNode
     *
     * @return
     */
    public JsonNode convertToReference() {
        ObjectNode node = JSONUtil.createObjectNode();
        node.put("type", "doc").put("fileName", this.getFileName())
                .put("fileId", this.getFileId());

        if (FileType.QA.getCode().equals(this.getDocType())) {
            if (StrUtil.isNotBlank(this.getAnswer())) {
                node.put("content", this.getContent() + "\n" + this.getAnswer());
            }
        } else {
            node.put("content", this.getContent());
        }

        return node;
    }

    public String getFileName() {
        return fileName;
    }

    public void setFileName(String fileName) {
        this.fileName = fileName;
    }

    public String getDocStatuesName() {
        return FileState.getByCode(getDocStatus()).getDescription();
    }


    public void setDocStatuesName(String docStatuesName) {
        this.docStatuesName = docStatuesName;
    }

    // Getters and Setters
    public Long getId() {
        return id;
    }

    public void setId(Long id) {
        this.id = id;
    }

    public String getContent() {
        return content;
    }

    public void setContent(String content) {
        this.content = content;
    }

    public Long getKnowledgeBaseId() {
        return knowledgeBaseId;
    }

    public void setKnowledgeBaseId(Long knowledgeBaseId) {
        this.knowledgeBaseId = knowledgeBaseId;
    }

    public Long getFileId() {
        return fileId;
    }

    public void setFileId(Long fileId) {
        this.fileId = fileId;
    }

    public String getDataType() {
        return dataType;
    }

    public void setDataType(String dataType) {
        this.dataType = dataType;
    }


    public Integer getDocStatus() {
        return docStatus;
    }

    public void setDocStatus(Integer docStatus) {
        this.docStatus = docStatus;
    }

    public Date getTrainTime() {
        return trainTime;
    }

    public void setTrainTime(Date trainTime) {
        this.trainTime = trainTime;
    }

    public Date getEmbedTime() {
        return embedTime;
    }

    public void setEmbedTime(Date embedTime) {
        this.embedTime = embedTime;
    }

    public Date getCreateTime() {
        return createTime;
    }

    public void setCreateTime(Date createTime) {
        this.createTime = createTime;
    }

    public String getAiKnowledgeBaseName() {
        return aiKnowledgeBaseName;
    }

    public void setAiKnowledgeBaseName(String aiKnowledgeBaseName) {
        this.aiKnowledgeBaseName = aiKnowledgeBaseName;
    }

    public String getDocType() {
        return docType;
    }

    public void setDocType(String docType) {
        this.docType = docType;
    }

    public Double getScore() {
        return score;
    }

    public void setScore(Double score) {
        this.score = score;
    }

    public String getLinkType() {
        return linkType;
    }

    public void setLinkType(String linkType) {
        this.linkType = linkType;
    }

    public Long getLinkId() {
        return linkId;
    }

    public void setLinkId(Long linkId) {
        this.linkId = linkId;
    }

    public String getAnswer() {
        return answer;
    }

    public void setAnswer(String answer) {
        this.answer = answer;
    }


    public Long getDocId() {
        return docId;
    }

    public void setDocId(Long docId) {
        this.docId = docId;
    }

    public Integer getQuestionNum() {
        return questionNum;
    }

    public void setQuestionNum(Integer questionNum) {
        this.questionNum = questionNum;
    }
}

